import atexit
import traceback
from typing import Type, TypeVar

import psycopg2
from psycopg2 import extras
from .DBDriver import DB
from config import yaml_config, DatabaseSetting, logger

T = TypeVar('T')

"""
    此类要求psycopg2的版本最低为2.9，因为下面两个语句的版本只针对于2.9以上版本
    from psycopg2 import extras
    extras.RealDictCursor
    若刚需psycopg2的2.9以下版本，则您需要修改RealDictCursor的引入方法
"""


class PostGreSQLDB(DB):
    def __init__(self):
        database_config: DatabaseSetting = yaml_config.database_setting
        super().__init__()
        self.connect: psycopg2.extensions.connection = psycopg2.connect(
            host=database_config.host,
            port=database_config.port,
            user=database_config.username,
            password=database_config.password,
            database=database_config.db,
            client_encoding=database_config.charset
        )
        self.cursor: psycopg2.extensions.cursor = self.connect.cursor(cursor_factory=extras.RealDictCursor)

    """
        获取数据库连接
    """

    def get_connect(self) -> psycopg2.extensions.connection:
        return self.connect

    """
        获取数据库游标
    """

    def get_cursor(self) -> psycopg2.extensions.cursor:
        return self.cursor

    """
        执行SQL语句，返回结果集
    """

    def execute(self, sql: str, params: tuple = None) -> list:
        try:
            self.cursor.execute(sql, params)
            if not self.transaction:
                self.connect.commit()
            result = self.cursor.fetchall()
            result_list: list = [x for x in result]
            return result_list
        except Exception as e:
            traceback.print_exc()
            self.connect.rollback()

    """
        执行SQL，获取一条数据，用对象接收
    """

    def get_one(self, sql: str, class_type: Type[T] = None, params: tuple = None) -> T:
        try:
            self.cursor.execute(sql, params)
            if not self.transaction:
                self.connect.commit()
            result: list = self.cursor.fetchone()
            obj = self.create_instance_from_type(result, class_type)
            return obj
        except Exception as e:
            traceback.print_exc()
            self.connect.rollback()

    """
        执行SQL，获取所有数据，用对象接收
    """

    def get_list(self, sql: str, class_type: Type[T] = None, params: tuple = None) -> list:
        self.cursor.execute(sql, params)
        if not self.transaction:
            self.connect.commit()
        result = self.cursor.fetchall()
        obj_list: list = []
        for row in result:
            obj = self.create_instance_from_type(row, class_type)
            obj_list.append(obj)
        return obj_list

    """
           执行一条insert的SQL，返回受影响行数
    """
    def insert_one(self, sql: str, params: tuple = None) -> int:
        self.cursor.execute(sql, params)
        row_count: int = self.cursor.rowcount
        if not self.transaction:
            self.connect.commit()
        return row_count

    """
        执行insert的SQL，插入多条数据，返回受影响行数
    """
    def insert_many(self, sql: str, params_list: list[tuple] = None) -> int:
        self.cursor.executemany(sql, params_list)
        row_count: int = self.cursor.rowcount
        if not self.transaction:
            self.connect.commit()
        return row_count
    """
        执行一条update的SQL，返回受影响行数
    """
    def update_one(self, sql: str, params: tuple = None) -> int:
        self.cursor.execute(sql, params)
        row_count: int = self.cursor.rowcount
        if not self.transaction:
            self.connect.commit()
        return row_count

    """
        执行一条delete的SQL，返回受影响行数
    """
    def delete(self, sql: str, params: tuple = None) -> int:
        self.cursor.execute(sql, params)
        row_count: int = self.cursor.rowcount
        if not self.transaction:
            self.connect.commit()
        return row_count

    """
        提交事务
    """
    def commit(self):
        try:
            self.connect.commit()
            logger.info(f"事务ID{self.transaction_id}提交成功")
        except Exception as e:
            traceback.print_exc()
            self.connect.rollback()
            logger.error(f"事务ID{self.transaction_id}提交失败")
        finally:
            transaction_id: int = self.transaction_id
            self.transaction = True
            self.transaction_id = None
            logger.info(f"事务ID{transaction_id}已关闭")

    def close(self):
        self.cursor.close()
        self.connect.close()
